# python3.7.1
# 这是一个水资源配置软件。

# 按 Shift+F10 执行或将其替换为您的代码。
# 按 双击 Shift 在所有地方搜索类、文件、工具窗口、操作和设置。

from reservoir import *
from construct import *
from const import *
import pickle
import traceback
import pandas as pd


class mainAllocation(object):
    # 水资源配置主程序，主要用于总处理各对外接口，及更新
    def __init__(self):
        # 首先建立基础框架
        self.reservoirlist = []  # 储水池
        self.constructlist = []  # 建筑

        # 全局参数区
        self.Step = 50  # 步数
        self.Steptime = 1  # 每一步的计算步长，单位是秒
        self.Relevel = 0.5

    def getserial(self, type):
        # 获取储水池/建筑物序列
        rserial = []
        for i in type:
            rserial.append(i.serial)
        return rserial

    def getname(self, type):
        # 获取储水池/建筑物序列
        namelist = []
        for i in type:
            namelist.append(i.name)
        return namelist

    def getreservoirname(self):
        # 获取储水池/建筑物序列
        reservoirnamelist = []
        for i in self.reservoirlist:
            reservoirnamelist.append(i.name)
        return reservoirnamelist

    def getreservoirserial(self):
        # 获取储水池/建筑物序列
        reservoirseriallist = []
        for i in self.reservoirlist:
            reservoirseriallist.append(i.serial)
        return reservoirseriallist

    def getconstructname(self):
        # 获取储水池/建筑物序列
        constructnamelist = []
        for i in self.constructlist:
            constructnamelist.append(i.name)
        return constructnamelist

    def getconstructserial(self):
        # 获取储水池/建筑物序列
        constructseriallist = []
        for i in self.constructlist:
            constructseriallist.append(i.serial)
        return constructseriallist

    def addreservoir(self, serial: float = 0, name: str = '储水池', firstlevel: float = 0,
                     upconstruct_serial: list = [],
                     downconstruct_serial: list = [], reservoirtype: str = 'reservoir',
                     levelvolumerelation: list = [[], []]):
        # 增加蓄水池
        try:
            # 为避免序号重复而进行检查
            if serial in self.getserial(self.reservoirlist):
                raise ValueError('添加失败，序号已存在，请更改')
            # 检查无误进行添加
            reservoir = Reservoir(serial=serial, name=name, firstlevel=firstlevel,
                                  upconstruct_serial=upconstruct_serial,
                                  downconstruct_serial=downconstruct_serial, reservoirtype=reservoirtype,
                                  levelvolumerelation=levelvolumerelation)
            self.reservoirlist.append(reservoir)
            print('增加储水池成功，名称为：' + reservoir.name)
            self._sortreservoirlist()
            return None
        except Exception as e:
            print('未能成功添加蓄水池', traceback.print_exc(), e)
            return e

    def addconstruct(self, serial, name, constructtype, upserial, downserial,
                     width, hight, maxflow=float('inf'), baselevel=0,
                     rule=None, target: list = None):
        # 增加建筑物
        if rule is None:
            rule = [float('inf'), float('-inf'), float('inf'), float('-inf')]
        try:
            # 为避免序号重复而进行检查
            if serial in self.getserial(self.constructlist):
                raise ValueError('序号已存在，请更改')
            if constructtype in ['双向闸', '单向闸']:
                construct = Construct(serial=serial, name=name, constructtype=constructtype, upserial=upserial,
                                      downserial=downserial, baselevel=baselevel, width=width, hight=hight,
                                      maxflow=maxflow, rule=rule)
            elif constructtype == '泵站':
                construct = pump(serial=serial, name=name, constructtype=constructtype, upserial=upserial,
                                 downserial=downserial, width=width, hight=hight, maxflow=maxflow,
                                 rule=rule, target=target)
            else:
                raise ValueError('建筑物类型有误')
            self.constructlist.append(construct)
            print('增加建筑物成功，名称为：' + construct.name)
            self._sortconstructlist()
            return None
        except Exception as e:
            print('未能成功添加建筑物', traceback.print_exc(), e)
            return e

    def delreservoir(self, serial):
        # 删除蓄水池    采用序号进行删除，采用名称将会导致重复

        # 判断是否存在该序号
        if serial in self.getserial(self.reservoirlist):
            for i in range(len(self.reservoirlist)):
                if self.reservoirlist[i].serial == serial:
                    del self.reservoirlist[i]
                    print('成功删除蓄水池', serial)
                    self._sortreservoirlist()
                    break
        else:
            print('未找到该序号的蓄水池，请检查')

    def delconstruct(self, serial):
        # 删除建筑物
        # 判断是否存在该序号
        if serial in self.getserial(self.constructlist):
            for i in range(len(self.constructlist)):
                if self.constructlist[i].serial == serial:
                    del self.constructlist[i]
                    print('成功删除建筑物', serial)
                    self._sortconstructlist()
                    break
        else:
            print('未找到该序号的蓄水池，请检查')
        pass

    def getreservoir(self, serial: float):
        # 获取当前序号的储水池
        for reservoir in self.reservoirlist:
            if reservoir.serial == serial:
                return reservoir
        raise ValueError('提供的序号错误，请检查')

    def getconstruct(self, serial):
        # 获取当前序号的建筑物
        for construct in self.constructlist:
            if construct.serial == serial:
                return construct
        raise ValueError('提供的序号错误，请检查')

    def _sortreservoirlist(self):
        # 根据序号调整储水池排序,已核验
        reservoirserial = self.getserial(self.reservoirlist)
        reservoirserial.sort()
        reservoirlist_temp = []
        for i in reservoirserial:
            for j in range(len(self.reservoirlist)):
                if self.reservoirlist[j].serial == i:
                    reservoirlist_temp.append(self.reservoirlist[j])
        if len(self.reservoirlist) == len(reservoirlist_temp):
            self.reservoirlist = reservoirlist_temp
            # print('已根据序号重新排序')
        else:
            raise ValueError('重新排序错误，请检查序号')

    def _sortconstructlist(self):
        # 根据序号调整建筑物排序,已核验
        constructserial = self.getserial(self.constructlist)
        constructserial.sort()
        constructlist_temp = []
        for i in constructserial:
            for j in range(len(self.constructlist)):
                if self.constructlist[j].serial == i:
                    constructlist_temp.append(self.constructlist[j])
        if len(self.constructlist) == len(constructlist_temp):
            self.constructlist = constructlist_temp
            # print('已根据序号重新排序')
        else:
            raise ValueError('重新排序错误，请检查序号')

    def allocation(self):
        if self.checkrelation():
            # 循环计算器
            for i in range(self.Step):
                try:
                    self.stepallocation()
                except Exception as e:
                    print('未能成功运算0', traceback.print_exc(), e)
        else:
            print('序号不匹配')

    def stepallocation(self):
        # 逐步计算

        # 以下是分别对建筑物、水库进行计算，但如果水库对应多个出流，则会出现流量溢出的问题，因此改以水库为主线，逐一计算
        # 首先根据上一时段末水位计算各建筑物过流
        for construct_now in self.constructlist:
            # 上一时段的上下游控制水位
            uplevel = self.getreservoir(construct_now.upserial).waterlevellist[-1]
            downlevel = self.getreservoir(construct_now.downserial).waterlevellist[-1]
            flow = construct_now.flow_calculation(uplevel, downlevel, self.Steptime)
            # 初步计算建筑物出流，再根据实际流量进行调整（避免出现负数）

        # 其次逐一对各水库水位进行调整
        for reservoir in self.reservoirlist:
            # 与之相连的建筑物入流情况
            inflow = 0
            outflow = 0
            for construct_in in reservoir.upconstruct_serial:
                inflow = inflow + self.getconstruct(construct_in).flowlist[-1]
            # print('inflow=', inflow)
            for construct_down in reservoir.downconstruct_serial:
                outflow = outflow + self.getconstruct(construct_down).flowlist[-1]
            # print('outflow=', outflow)
            flow = inflow - outflow

            # 验证当前入流是否满足
            reoutflow = reservoir.waterchange(flow)  # 当水量可供调蓄时，返回值应为0，当水量不足调蓄，返回差值(负数）
            if reoutflow != 0:
                # 水量不足，即根据排序，自前向后调整。
                for construct_down in reservoir.downconstruct_serial[::-1]:
                    # 逆序遍历
                    reoutflow = self.getconstruct(construct_down).reflow(reoutflow)
                if reoutflow > 0:
                    raise ValueError('水量无非平衡')
                # reoutflow = reoutflow / len(reservoir.downconstruct_serial)  # 平均分配水量，仅调整下游建筑物
                # for construct_down in reservoir.downconstruct_serial:
                #     self.getconstruct(construct_down).flowlist[-1] = (reoutflow +
                #                                                       self.getconstruct(construct_down).flowlist[-1])

    def reflash(self):
        # 删除已计算的值：蓄水池的 waterlevellist watervolumelist
        for reservoir in self.reservoirlist:
            reservoir.waterlevellist = []  # 给出水位序列
            reservoir.watervolumelist = []  # 给出库容序列
            reservoir.waterlevellist.append(reservoir.firstlevel)  # 初始水位
            if reservoir.reservoirtype == 'reservoir':
                reservoir.watervolumelist.append(reservoir.leveltovolume(reservoir.firstlevel))  # 初始库容
            else:
                reservoir.watervolumelist = [float('inf')]  # 当为河流时，库容为无穷大
        for construct in self.constructlist:
            construct.statelist = []
            construct.flowlist = []

    def checkrelation(self):
        # 检查储水池与建筑物的关系
        for reservoir in self.reservoirlist:
            # 首先查找储水池
            for upconstruct in reservoir.upconstruct_serial:
                # 检查上游建筑物的出口建筑物序号，是否与建筑物相对应
                if upconstruct:
                    try:
                        construct = self.getconstruct(upconstruct)
                        if construct.downserial == reservoir.serial:
                            print('OK')
                        else:
                            print('当前储水池%s，与上游建筑物%s不匹配' % (reservoir.serial, construct.serial))
                            return False
                    except Exception as e:
                        print(e)
                        return False
                else:
                    print('无上游建筑物')
            for downconstruct in reservoir.downconstruct_serial:
                # 检查下游建筑物的进口建筑物序号，是否与建筑物相对应
                if downconstruct:
                    try:
                        construct = self.getconstruct(downconstruct)
                        if construct.upserial == reservoir.serial:
                            print('OK')
                        else:
                            print('当前储水池%s，与下游建筑物%s不匹配' % (reservoir.serial, construct.serial))
                            return False
                    except Exception as e:
                        print(e)
                        return False
                else:
                    print('无下游建筑物')
        for construct in self.constructlist:
            # 检查建筑物的上下游储水池的上下游接建筑物是否包含本建筑物
            try:
                upreservoir = self.getreservoir(construct.upserial)
                downreservoir = self.getreservoir(construct.downserial)
            except Exception as e:
                print(e)
                return False
            if construct.serial in upreservoir.downconstruct_serial:
                # 建筑物上游储水池的下游序列，应当包含本建筑物
                print('建筑物与上游储水池相对应')
            else:
                print('%s建筑物与上游储水池%s不对应' % (construct.serial, upreservoir.serial))
                return False
            if construct.serial in downreservoir.upconstruct_serial:
                # 建筑物下游储水池的上游序列，应当包含本建筑物
                print('建筑物与下游储水池相对应')
            else:
                print('%s建筑物与下游储水池%s不对应' % (construct.serial, downreservoir.serial))
                print(construct.downserial, downreservoir.upconstruct_serial)
                return False
        return True

    def saveExcel(self, serial=[], name=[], firstlevel=[], upconstruct_serial=[], downconstruct_serial=[],
                  reservoirtype=[], levelvolumerelation=[]):
        # 将相关参数都保存到excel
        io = r'test.xlsx'  # 相对路径

        # 全局设置
        globalset_df = pd.DataFrame({
            '时间步长（s）': self.Steptime,
            '步数': self.Step,
            '水位要求': self.Relevel
        }, index=['全局设置'])
        # globalset = pd.DataFrame({
        #     self.Steptime,
        #     self.Step,
        #     self.Relevel
        # }, index=['时间步长（s）', '步数', '水位要求'])
        # globalset_df.to_excel(io, sheet_name='全局设置', index=True)

        # 储水池设置
        # serial, name, firstlevel, upconstruct_serial, downconstruct_serial, type, levelvolumerelation = [[]*7]
        for reservoir in self.reservoirlist:
            serial.append(reservoir.serial)
            name.append(reservoir.name)
            firstlevel.append(reservoir.firstlevel)
            upconstruct_serial.append(reservoir.upconstruct_serial)
            downconstruct_serial.append(reservoir.downconstruct_serial)
            reservoirtype.append(reservoir.reservoirtype)
            levelvolumerelation.append(reservoir.levelvolumerelation)
        print('serial', serial)
        reservoir_df = pd.DataFrame({
            '序号': serial,
            '名称': name,
            '初设水位': firstlevel,
            '上游建筑物序列': upconstruct_serial,
            '下游建筑物序列': downconstruct_serial,
            '储水池类型': reservoirtype,
            '水位流量关系': levelvolumerelation
        }, index=range(len(serial)))
        # print(reservoir_df)
        reservoir_df.to_excel(io, sheet_name='储水池', index=True)
        pass

    def readExcel(self):
        # 从excel中读取数据
        pass


def main():
    # 建立总库
    mainallocation = mainAllocation()
    # 增加储水池与建筑物
    mainallocation.addreservoir(serial=0, name='水库1', firstlevel=100, downconstruct_serial=[0, 1, 3, 4],
                                reservoirtype='reservoir', levelvolumerelation=levelvolumerelation0)
    mainallocation.addreservoir(serial=1, name='河流1', firstlevel=50, upconstruct_serial=[0],
                                reservoirtype='river')  # 河流
    mainallocation.addreservoir(serial=2, name='水库2', firstlevel=70, upconstruct_serial=[1],
                                reservoirtype='river')  # 河流
    mainallocation.addreservoir(serial=1.2, name='水库3', firstlevel=70, upconstruct_serial=[3, 4],
                                reservoirtype='river')  # 河流
    mainallocation.addconstruct(serial=0, name='建筑物1', constructtype='双向闸', upserial=0, downserial=1.0,
                                baselevel=0, width=1, hight=1, maxflow=1000)
    mainallocation.addconstruct(serial=1, name='建筑物2', constructtype='单向闸', upserial=0, downserial=2.0,
                                baselevel=0, width=1, hight=1, maxflow=100)
    mainallocation.addconstruct(serial=3, name='建筑物3', constructtype='单向闸', upserial=0, downserial=1.2,
                                baselevel=0, width=1, hight=1, maxflow=100)
    mainallocation.addconstruct(serial=4, name='建筑物4', constructtype='泵站', upserial=0, downserial=1.2, width=1,
                                hight=1, maxflow=100)

    print('储水池序号', mainallocation.getserial(mainallocation.reservoirlist),
          mainallocation.getname(mainallocation.reservoirlist))
    # print(mainallocation.getserial(mainallocation.constructlist))
    serial = 1.2
    print(mainallocation.getreservoir(serial).name)

    print('储水池序号', mainallocation.getserial(mainallocation.reservoirlist),
          mainallocation.getname(mainallocation.reservoirlist))

    print(mainallocation.getserial(mainallocation.constructlist))

    mainallocation.allocation()  # 运算50步

    # 查看建筑物水位过程
    for reservoir in mainallocation.reservoirlist:
        print('第%s号水库的水位变化情况' % reservoir.serial)
        print(reservoir.waterlevellist)
        print('第%s号水库的库容变化情况' % reservoir.serial)
        print(reservoir.watervolumelist)

    for construct_now in mainallocation.constructlist:
        print('第%s号建筑物的流量变化情况' % construct_now.serial)
        print(construct_now.flowlist)

    # mainallocation.reflash()

    # # 查看建筑物水位过程
    # for reservoir in mainallocation.reservoirlist:
    #     print('第%s号水库的水位变化情况' % reservoir.serial)
    #     print(reservoir.waterlevellist)
    #     print('第%s号水库的库容变化情况' % reservoir.serial)
    #     print(reservoir.watervolumelist)
    #
    # for construct_now in mainallocation.constructlist:
    #     print('第%s号建筑物的流量变化情况' % construct_now.serial)
    #     print(construct_now.flowlist)

    # 保存
    # with open('test.pickle', 'wb') as file:
    #     pickle.dump(mainallocation, file)
    #     print('保存成功')
    # with open('test.pickle', 'rb') as file:
    #     abc = pickle.load(file)
    #     print('读取成功')
    #     print(abc.getreservoirname())

    # EXCEL 相关操作
    mainallocation.saveExcel()


# 按装订区域中的绿色按钮以运行脚本。
if __name__ == '__main__':
    main()

# 访问 https://www.jetbrains.com/help/pycharm/ 获取 PyCharm 帮助
